from pylab import *
from numpy import * #unnecessary, just for pydev auto-complete


class MatchingPursuit():
    def __init__(self, length, startingFreq, numFreq, numFreqPerOctave):
        self.init_basis(length, startingFreq, numFreq, numFreqPerOctave)
        
    def init_basis(self, length, startingFreq, numFreq, numFreqPerOctave):
        self.frequency = startingFreq * exp(log(2) * arange(0, numFreq) / numFreqPerOctave)
        self.angular_frequency = 2 * pi * self.frequency
        self.basis_sin = matrix(sin(outer(self.angular_frequency, linspace(0, 1, length, False))))
        self.basis_cos = matrix(cos(outer(self.angular_frequency, linspace(0, 1, length, False))))
        
    def matching_pursuit(self, data):
        c = self.basis_cos * data.reshape(-1, 1) / data.size
        s = self.basis_sin * data.reshape(-1, 1) / data.size
        tmp = sqrt(power(c, 2) + power(s, 2)) * 2
#        plot(tmp)
#        show()
        I = tmp.argmax()
        A = tmp.item(I)
        w = self.angular_frequency[I]
        p = arctan(c.item(I) / s.item(I)) - pi/2
        if s.item(I) < 0:
            p += pi
        return array([A, w, p])

if __name__ == '__main__':
    a = MatchingPursuit(44100 / 10, 50, 120 * 10, 120)
    data = random.random((44100 / 10, 1))
    data = ones((44100 / 10, 1))
    data = array(20 * cos(100 * 2 * pi * linspace(0, 1, 44100 / 10, False).reshape(-1, 1)))
    [A, w, p] = a.matching_pursuit(data)
    
    plot(A * cos(w * linspace(0, 1, 44100 / 10, False) + p))
    plot(data)
    show()
